/*
 * PyRSA.c
 *
 *  Created on: 2009-5-9
 *      Author: thh
 */


#include "base.h"

static PyMethodDef method[] = {
		{"size",(PyCFunction)PyRSA_Size,METH_NOARGS,"get rsa size"},
		{"check",(PyCFunction)PyRSA_Check,METH_NOARGS,"check rsa key pair"},
		{"public_encrypt",(PyCFunction)PyRSA_public_encrypt,METH_VARARGS,""},
		{"public_decrypt",(PyCFunction)PyRSA_public_decrypt,METH_VARARGS,""},
		{"private_decrypt",(PyCFunction)PyRSA_private_decrypt,METH_VARARGS,""},
		{"private_encrypt",(PyCFunction)PyRSA_private_encrypt,METH_VARARGS,""},
		{"sign",(PyCFunction)PyRSA_sign,METH_VARARGS,""},
		{"verify",(PyCFunction)PyRSA_verify,METH_VARARGS,""},
		{NULL}
};

PyTypeObject PyRSA_Type = {
		PyObject_HEAD_INIT(NULL)
		//{{0,0,1},0},
		MODULE_NAME".RSA",	//const char *tp_name; /* For printing, in format "<module>.<name>" */
		sizeof(PyRSA),0, 	//Py_ssize_t tp_basicsize, tp_itemsize; /* For allocation */

		/* Methods to implement standard operations */

		(destructor)PyRSA_Del,			//destructor tp_dealloc;
		NULL,			//printfunc tp_print;
		NULL,			//getattrfunc tp_getattr;
		NULL,			//setattrfunc tp_setattr;
		NULL,			//void *tp_reserved; /* formerly known as tp_compare */
		NULL,			//reprfunc tp_repr;

		/* Method suites for standard classes */

		NULL,			//PyNumberMethods *tp_as_number;
		NULL,			//PySequenceMethods *tp_as_sequence;
		NULL,			//PyMappingMethods *tp_as_mapping;

		/* More standard operations (here for binary compatibility) */

		NULL,			//hashfunc tp_hash;
		NULL,			//ternaryfunc tp_call;
		(reprfunc)PyRSA_Str,			//reprfunc tp_str;
		NULL,			//getattrofunc tp_getattro;
		NULL,			//setattrofunc tp_setattro;

		/* Functions to access object as input/output buffer */
		NULL,			//PyBufferProcs *tp_as_buffer;

		/* Flags to define presence of optional/expanded features */
		Py_TPFLAGS_DEFAULT, //long tp_flags;

		"RSA object",//const char *tp_doc; /* Documentation string */

		/* Assigned meaning in release 2.0 */
		/* call function for all accessible objects */
		NULL,			//traverseproc tp_traverse;

		/* delete references to contained objects */
		NULL,			//inquiry tp_clear;

		/* Assigned meaning in release 2.1 */
		/* rich comparisons */
		NULL,			//richcmpfunc tp_richcompare;

		/* weak reference enabler */
		0,			//Py_ssize_t tp_weaklistoffset;

		/* Iterators */
		NULL,			//getiterfunc tp_iter;
		NULL,			//iternextfunc tp_iternext;

		/* Attribute descriptor and subclassing stuff */
		method,			//struct PyMethodDef *tp_methods;
		NULL,			//struct PyMemberDef *tp_members;
		NULL,			//struct PyGetSetDef *tp_getset;
		NULL,			//struct _typeobject *tp_base;
		NULL,			//PyObject *tp_dict;
		NULL,			//descrgetfunc tp_descr_get;
		NULL,			//descrsetfunc tp_descr_set;
		0,			//Py_ssize_t tp_dictoffset;
		(initproc)PyRSA_Init,			//initproc tp_init;
		NULL,			//allocfunc tp_alloc;
		(newfunc)PyType_GenericNew,			//newfunc tp_new;
		NULL,			//freefunc tp_free; /* Low-level free-memory routine */
		NULL,			//inquiry tp_is_gc; /* For PyObject_IS_GC */
		NULL,			//PyObject *tp_bases;
		NULL,			//PyObject *tp_mro; /* method resolution order */
		NULL,			//PyObject *tp_cache;
		NULL,			//PyObject *tp_subclasses;
		NULL,			//PyObject *tp_weaklist;
		NULL,			//destructor tp_del;

		/* Type attribute cache version tag. Added in version 2.6 */
		0			//unsigned int tp_version_tag;
};


int PyRSA_Init(PyRSA *self,PyObject *args)
{
	unsigned long e = 0;
	int bits = 0;

	if(!PyArg_ParseTuple(args,"il",&bits,&e)){
		return -1;
	}
	if( e%2 != 1 ){
		PyErr_SetString(PyExc_RuntimeError,"e must be an odd number");
		return -1;
	}
	self->rsa = RSA_generate_key(bits,e,NULL,NULL);
	if(self->rsa==NULL){
		PyErr_SetString(PyExc_RuntimeError,"generate rsa key failed");
		return -1;
	}
	return 1;
}

void PyRSA_Del(PyRSA *self)
{
	if(self->rsa!=NULL)
		RSA_free(self->rsa);
	PyObject_DEL(self);
}

PyObject *PyRSA_Str(PyRSA *self)
{
	PyObject *retval;
	char *txt;
	long len;

	BIO *mem = BIO_new(BIO_s_mem());
	if(!RSA_print(mem,self->rsa,0)){
		PyErr_SetString(PyExc_RuntimeError,"print rsa failed");
		return NULL;
	}
	len = BIO_get_mem_data(mem,&txt);
	retval = PyUnicode_FromStringAndSize(txt,len);
	BIO_free(mem);
	return retval;
}

PyObject *PyRSA_Size(PyRSA *self)
{
	return Py_BuildValue("i",RSA_size(self->rsa));
}

PyObject *PyRSA_Check(PyRSA *self)
{
	int v = RSA_check_key(self->rsa);
	switch(v){
	case 1:
		return Py_True;
	case 0:
		return Py_False;
	default:
		PyErr_SetString(PyExc_RuntimeError,"check rsa failed");
		return NULL;
	}
}

typedef int	(*enc_dec_func)(int flen, const unsigned char *from,
		unsigned char *to, RSA *rsa,int padding);

PyObject *PyRSA_enc_dec_func(PyRSA *self,PyObject *args,enc_dec_func func)
{
	PyObject *retval;
	unsigned char *buf = NULL,*to;
	int len = 0, padding=RSA_PKCS1_PADDING;
	if(!PyArg_ParseTuple(args,"y#|i",&buf,&len,&padding))
		return NULL;

	to = PyMem_Malloc(RSA_size(self->rsa));
	len = (*func)(len,buf,to,self->rsa,padding);
	if(len==-1){
		PyMem_Free(to);
		PyErr_OpenSSLError();
		return NULL;
	}
	retval = Py_BuildValue("y#",to,len);
	PyMem_Free(to);
	return retval;
}

PyObject *PyRSA_public_encrypt(PyRSA *self,PyObject *args)
{
	return PyRSA_enc_dec_func(self,args,RSA_public_encrypt);
}

PyObject *PyRSA_private_decrypt(PyRSA *self,PyObject *args)
{
	return PyRSA_enc_dec_func(self,args,RSA_private_decrypt);
}

PyObject *PyRSA_public_decrypt(PyRSA *self,PyObject *args)
{
	return PyRSA_enc_dec_func(self,args,RSA_public_decrypt);
}

PyObject *PyRSA_private_encrypt(PyRSA *self,PyObject *args)
{
	return PyRSA_enc_dec_func(self,args,RSA_private_encrypt);
}

PyObject *PyRSA_sign(PyRSA *self,PyObject *args)
{
	PyObject *retval;
	unsigned char *msg,*sigret;
	unsigned int msglen = 0, keylen=RSA_size(self->rsa);
	unsigned int type = NID_sha1;
	if(!PyArg_ParseTuple(args,"y#|i",&msg,&msglen,&type))
		return NULL;

	sigret = PyMem_MALLOC(keylen);
	if(RSA_sign(type,msg,msglen,sigret,&keylen,self->rsa)!=1){
		PyMem_Free(sigret);
		PyErr_OpenSSLError();
		return NULL;
	}
	retval = Py_BuildValue("y#",sigret,keylen);
	PyMem_Free(sigret);
	return retval;
}

PyObject *PyRSA_verify(PyRSA *self,PyObject *args)
{
	unsigned char *msg,*sigret;
	int msglen = 0, keylen=RSA_size(self->rsa);
	int type = NID_sha1;
	if(!PyArg_ParseTuple(args,"y#y#|i",&msg,&msglen,&sigret,&keylen,&type))
		return NULL;
	int retval = RSA_verify(type,msg,msglen,sigret,keylen,self->rsa);
	switch(retval){
		case 1: return Py_True;
		case 0: return Py_False;
		default:
			PyErr_OpenSSLError();
			return NULL;
	}
}
